package hw2;

import java.util.HashMap;
import java.util.Iterator;

public class BigramModelLearner
{
	HashMap<Character, HashMap<Character, Integer>> model;
	
	public BigramModelLearner()
	{
		model = new HashMap<Character, HashMap<Character, Integer>>();
	}
	
	public void train(String trainingInput)
	{
		int characterIndex = 0;
		char previousCharacter = 0;
		while (characterIndex < trainingInput.length())
		{
			char currentCharacter = trainingInput.charAt(characterIndex);
			
			if (model.get(previousCharacter) == null)
				model.put(previousCharacter, new HashMap<Character, Integer>());
			
			if (model.get(previousCharacter).get(currentCharacter) == null)
				model.get(previousCharacter).put(currentCharacter, 0);
			
			model.get(previousCharacter).put(currentCharacter, 
					model.get(previousCharacter).get(currentCharacter) + 1);
			
			characterIndex++;
			previousCharacter = currentCharacter;
		}
		
		if (model.get(previousCharacter) == null)
			model.put(previousCharacter, new HashMap<Character, Integer>());
	}
	
	public void test(String testInput)
	{
		int numCorrectPredictions = 0;
		
		int characterIndex = 0;
		char previousCharacter = 0;
		while (characterIndex < testInput.length())
		{
			// Predict next character.
			char predictedCharacter = nextPrediction(previousCharacter);
			
			// Read character.
			char actualNextCharacter = testInput.charAt(characterIndex);
			
			// Determine if prediction was correct.
			if (predictedCharacter == actualNextCharacter)
				numCorrectPredictions++;
			
			// Update model.
			if (model.get(actualNextCharacter) == null)
				model.put(actualNextCharacter, new HashMap<Character, Integer>());
			
			if (model.get(previousCharacter).get(actualNextCharacter) == null)
				model.get(previousCharacter).put(actualNextCharacter, 0);
			
			model.get(previousCharacter).put(actualNextCharacter, 
					model.get(previousCharacter).get(actualNextCharacter) + 1);
			
			characterIndex++;
			previousCharacter = actualNextCharacter;
		}
		
		// Print percentage of correct predictions.
		System.out.println("Percentage of correct predictions: " + 
				((numCorrectPredictions / (float) testInput.length()) * 100) + "%"); 
		
		// Print random prediction performance (1 / number of unique symbols).
		System.out.println("Random prediction performance: " + 
				(1 / (float) (model.keySet().size() - 1)) * 100 + "%");
	}
	
	private char nextPrediction(char previousCharacter)
	{
		int highestFrequency = 0;
		char mostFrequentCharacter = 0;
		
		if (model.get(previousCharacter) == null)
			return model.keySet().iterator().next();
		
		Iterator<Character> modelIterator = 
			model.get(previousCharacter).keySet().iterator();
		
		while (modelIterator.hasNext())
		{
			char nextCharacter = modelIterator.next();
			int nextFrequency = model.get(previousCharacter).get(nextCharacter);
			if (nextFrequency > highestFrequency)
			{
				highestFrequency = nextFrequency;
				mostFrequentCharacter = nextCharacter;
			}
		}
		
		return mostFrequentCharacter;
	}
}
